package info.wangyuan.agent.service.conversation;


import cn.hutool.core.collection.CollUtil;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import dev.langchain4j.data.message.SystemMessage;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.TokenStream;
import info.wangyuan.agent.common.Constant.AgentConstant;
import info.wangyuan.agent.common.Constant.PromptTemplate;
import info.wangyuan.agent.common.Enum.MessageType;
import info.wangyuan.agent.common.Enum.MsgEnum;
import info.wangyuan.agent.config.ModelFactory;
import info.wangyuan.agent.config.RedisChatMemoryStore;
import info.wangyuan.agent.entity.dto.AgentChatResponseDTO;
import info.wangyuan.agent.entity.dto.RagStreamChatRequestDTO;
import info.wangyuan.agent.entity.dto.RetrievedDocumentDTO;
import info.wangyuan.agent.entity.po.*;
import info.wangyuan.agent.exception.BusinessException;
import info.wangyuan.agent.mapper.AgentConfigMapper;
import info.wangyuan.agent.mapper.FileInfoMapper;
import info.wangyuan.agent.mapper.KnowledgeSourceMapper;
import info.wangyuan.agent.mapper.PromptMapper;
import info.wangyuan.agent.service.curd.KnowledgeBaseService;
import info.wangyuan.agent.service.rag.RagEmbeddingDomainService;
import info.wangyuan.agent.utils.JsonUtils;
import info.wangyuan.agent.utils.SseEmitterUtils;
import info.wangyuan.agent.utils.WordsFilter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Component;
import org.springframework.util.StopWatch;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeoutException;
import java.util.stream.Collectors;

/**
 * 流式对话处理流程编排
 *
 * @author Albert
 * @since 2025-09-07 21:16:53
 */
@Slf4j
@Component
public class ConversationDomainService {
    private final ModelFactory modelFactory;
    private final PromptMapper promptMapper;
    private final FileInfoMapper fileInfoMapper;
    private final AgentConfigMapper agentConfigMapper;
    private final ThreadPoolTaskExecutor ttlAsyncExecutor;
    private final KnowledgeBaseService knowledgeBaseService;
    private final RedisChatMemoryStore redisChatMemoryStore;
    private final KnowledgeSourceMapper knowledgeSourceMapper;
    private final RagEmbeddingDomainService ragEmbeddingDomainService;

    public ConversationDomainService(ModelFactory modelFactory, PromptMapper promptMapper, FileInfoMapper fileInfoMapper, AgentConfigMapper agentConfigMapper, ThreadPoolTaskExecutor ttlAsyncExecutor, KnowledgeBaseService knowledgeBaseService, RedisChatMemoryStore redisChatMemoryStore, KnowledgeSourceMapper knowledgeSourceMapper, RagEmbeddingDomainService ragEmbeddingDomainService) {
        this.modelFactory = modelFactory;
        this.promptMapper = promptMapper;
        this.fileInfoMapper = fileInfoMapper;
        this.agentConfigMapper = agentConfigMapper;
        this.ttlAsyncExecutor = ttlAsyncExecutor;
        this.knowledgeBaseService = knowledgeBaseService;
        this.redisChatMemoryStore = redisChatMemoryStore;
        this.knowledgeSourceMapper = knowledgeSourceMapper;
        this.ragEmbeddingDomainService = ragEmbeddingDomainService;
    }

    /**
     * RAG流式问答
     *
     * @param request 流式问答请求
     * @return SSE流式响应
     */
    public SseEmitter ragStreamChat(RagStreamChatRequestDTO request) {
        SseEmitter emitter = new SseEmitter(Long.MAX_VALUE);

        String sessionId = request.getSessionId();

        // 对用户输入进行内容安全校验，防止越狱攻击和角色扮演提示词，引导模型生成不良内容
        if (WordsFilter.wordsCheck(request.getQuestion())) {
            sendSseData(emitter, createErrorResponse(MsgEnum.WORDS_CHECK_ERROR.getMessage()));
            return emitter;
        }

        // 设置连接关闭回调
        emitter.onCompletion(() -> log.info("RAG 流式响应完毕: sessionId={}", sessionId));
        emitter.onTimeout(() -> {
            log.warn("RAG 流式响应超时: sessionId={}", sessionId);
            sendSseData(emitter, createErrorResponse("连接超时"));
        });
        emitter.onError((ex) -> {
            log.error("RAG 流式响应连接错误: sessionId={}, error={}", sessionId, ex.getMessage(), ex);
            sendSseData(emitter, createErrorResponse("服务暂时不可用，请稍后再试"));
        });

        CompletableFuture.runAsync(
                        () -> processRagStreamChat(request, emitter),
                        ttlAsyncExecutor
                )
                .exceptionally(ex -> {
                    log.error("RAG 流式响应异常: sessionId={}, error={}", request.getSessionId(), ex.getMessage(), ex);
                    sendSseData(emitter, createErrorResponse("服务暂时不可用，请稍后再试"));
                    SseEmitterUtils.safeComplete(emitter);
                    return null;
                });

        return emitter;
    }

    /**
     * 处理RAG流式问答的核心逻辑
     */
    private void processRagStreamChat(RagStreamChatRequestDTO request, SseEmitter emitter) {
        try {
            // 第一阶段：检索文档
            log.info("开始 RAG 流式聊天，问题：'{}'", request.getQuestion());

            // 发送检索开始信号
            sendSseData(emitter, AgentChatResponseDTO.build("开始检索相关文档...", MessageType.RAG_RETRIEVAL_START));
//            Thread.sleep(500);

            // 确定检索范围
            AgentConfig agentConfig = agentConfigMapper.selectOne(
                    Wrappers.lambdaQuery(AgentConfig.class)
                            .eq(AgentConfig::getEnabled, AgentConstant.ON)
            );
            if (agentConfig == null) {
                log.error("未找到启用的 Agent 配置");
                sendSseData(emitter, createErrorResponse("服务暂时不可用，请稍后再试"));
                throw new BusinessException("未找到启动的 Agent 配置");
            }

            Integer activeAgentId = agentConfig.getId();

            List<Integer> searchDatasetIds = knowledgeSourceMapper.selectList(
                            Wrappers.lambdaQuery(KnowledgeSource.class)
                                    .eq(KnowledgeSource::getAgentId, activeAgentId)
                    )
                    .stream()
                    .map(KnowledgeSource::getKnowledgeBaseId)
                    .toList();

            // 检查知识库是否存在
            knowledgeBaseService.checkKnowledgeBasesExist(searchDatasetIds);

            sendSseData(emitter, AgentChatResponseDTO.build("正在知识库中检索...", MessageType.RAG_RETRIEVAL_PROGRESS));

            // 执行RAG检索
            List<DocumentUnit> retrievedDocuments = ragEmbeddingDomainService.ragDoc(
                    searchDatasetIds,
                    request.getQuestion(),
                    request.getMaxResults(),
                    request.getMinScore(),
                    request.getEnableRerank(),
                    2,
                    Boolean.FALSE
            ); // 流式问答中暂时不启用查询扩展，保持现有行为

            // 构建检索结果
            List<RetrievedDocumentDTO> retrievedDocs = new ArrayList<>();
            List<Integer> fileIds = retrievedDocuments.stream()
                    .map(DocumentUnit::getFileId)
                    .distinct()
                    .toList();

            Map<Integer, String> fileNameMap = CollUtil.isEmpty(fileIds) ? Map.of() :
                    fileInfoMapper.selectByIds(fileIds).stream()
                            .collect(Collectors.toMap(FileInfo::getId, FileInfo::getName));

            for (DocumentUnit doc : retrievedDocuments) {
                double similarityScore = doc.getSimilarityScore() != null ? doc.getSimilarityScore() : 0.0;
                String fileName = fileNameMap.getOrDefault(doc.getFileId(), "未知文件");
                retrievedDocs.add(new RetrievedDocumentDTO(doc.getFileId(), fileName, doc.getId(), similarityScore));
            }

            // 发送检索完成信号
            String retrievalMessage = String.format("检索完成，找到 %d 个相关文档", retrievedDocs.size());
            AgentChatResponseDTO retrievalEndResponse = AgentChatResponseDTO.build(retrievalMessage, MessageType.RAG_RETRIEVAL_END);
            retrievalEndResponse.setPayload(JsonUtils.toJsonString(retrievedDocs));
            sendSseData(emitter, retrievalEndResponse);
//            Thread.sleep(500);

            // 第二阶段：生成回答
            sendSseData(emitter, AgentChatResponseDTO.build("开始生成回答...", MessageType.RAG_ANSWER_START));
//            Thread.sleep(500);

            // 构建LLM上下文
            String context = buildContextFromDocuments(retrievedDocuments);
            String prompt = buildRagPrompt(request.getQuestion(), context);
            log.info("RAG 提示：{}", prompt);
            // 调用流式LLM - 使用同步等待确保流式处理完成
            generateStreamAnswerAndWait(prompt, emitter, request.getSessionId());

        } catch (Exception e) {
            log.error("RAG流式聊天处理中出现错误", e);
            sendSseData(emitter, createErrorResponse("服务暂时不可用，请稍后再试"));
        }
    }

    /**
     * 发送SSE数据（带状态检查）
     */
    private void sendSseData(SseEmitter emitter, AgentChatResponseDTO response) {
        try {
            String jsonData = JsonUtils.toJsonString(response);
            SseEmitterUtils.safeSend(emitter, jsonData);
        } catch (Exception e) {
            log.error("发送SSE数据失败", e);
        }
    }

    /**
     * 构建检索文档的上下文
     */
    private String buildContextFromDocuments(List<DocumentUnit> documents) {
        if (CollUtil.isEmpty(documents)) return "暂无相关文档信息。";

        // 提取所有 fileId
        List<Integer> fileIds = documents.stream()
                .map(DocumentUnit::getFileId)
                .distinct()
                .toList();

        // 一次性查询并映射成 Map
        Map<Integer, String> fileNameMap = CollUtil.isEmpty(fileIds) ? Map.of() :
                fileInfoMapper.selectByIds(fileIds).stream()
                        .collect(Collectors.toMap(FileInfo::getId, FileInfo::getName));

        StringBuilder context = new StringBuilder("以下是相关的文档片段：\n\n");

        for (DocumentUnit doc : documents) {
            String fileName = fileNameMap.getOrDefault(doc.getFileId(), "未知文件");
            int page = doc.getPageIndex() + 1; // 优先用真实页码
            context.append(String.format("%s 第 %d 页：\n", fileName, page));
            context.append(doc.getContent());
            context.append("\n\n");
        }

        return context.toString();
    }

    /**
     * 构建RAG提示词
     */
    private String buildRagPrompt(String question, String context) {
        return String.format(PromptTemplate.RAG_TO_LLM_PROMPT, context, question);
    }

    /**
     * 创建错误响应
     */
    private AgentChatResponseDTO createErrorResponse(String errorMessage) {
        return AgentChatResponseDTO.buildEndMessage(errorMessage, MessageType.TEXT);
    }

    /**
     * 生成流式回答
     *
     * @param prompt    RAG提示词
     * @param emitter   SSE连接
     * @param sessionId 会话ID
     */
    private void generateStreamAnswerAndWait(String prompt, SseEmitter emitter, String sessionId) {
        try {
            StopWatch stopWatch = new StopWatch("RAG回答生成...");
            log.info("开始生成RAG回答");
            stopWatch.start("streaming-process");

            // 创建流式LLM客户端
            StreamingChatModel streamingChatModel = modelFactory.getCustomStreamingChatModel();

            // 创建Agent并启动流式处理
            Agent agent = buildStreamingAgent(streamingChatModel, sessionId);

            TokenStream tokenStream = agent.chat(prompt);

            // 使用CompletableFuture来等待流式处理完成
            CompletableFuture<Void> streamComplete = new CompletableFuture<>();

            // 累积思维链内容 --> 兼容有思考过程的 LLM
            final StringBuilder thinkingBuffer = new StringBuilder();
            final boolean[] thinkingStarted = {false};

            tokenStream.onPartialThinking(partial -> {
                String thought = partial.text();

                // 第一次收到思考片段 --> 推送“开始思考”
                if (!thinkingStarted[0]) {
                    thinkingStarted[0] = true;
                    sendSseData(emitter, AgentChatResponseDTO.build("开始思考...", MessageType.RAG_THINKING_START));
                }

                // 累积
                thinkingBuffer.append(thought);

                // 日志记录思考片段
                log.debug("收到思考片段: {}", thought);

                // 推送给前端
                sendSseData(emitter, AgentChatResponseDTO.build(thought, MessageType.RAG_THINKING_PROGRESS));
            }).onPartialResponse(fragment -> {

                log.debug("收到回答片段: {}", fragment);
                sendSseData(emitter, AgentChatResponseDTO.build(fragment, MessageType.RAG_ANSWER_PROGRESS));
            }).onCompleteResponse(chatResponse -> {
                // 如果有思考过程 --> 打印完整日志并结束
                if (!thinkingBuffer.isEmpty()) {
                    log.info("思考完成");
                    log.info("完整思维链内容:\n{}", thinkingBuffer);
                    sendSseData(emitter, AgentChatResponseDTO.build("思考完成", MessageType.RAG_THINKING_END));
                }

                String fullAnswer = chatResponse.aiMessage() != null ? chatResponse.aiMessage().text() : null;
                if (fullAnswer != null) {
                    log.info("RAG回答生成完成，响应长度: {}", fullAnswer.length());
                    log.info("完整RAG回答内容:\n{}", fullAnswer);
                } else {
                    log.info("RAG回答生成完成，无内容");
                }

                // 告诉前端完成
                sendSseData(emitter, AgentChatResponseDTO.buildEndMessage("回答生成完成", MessageType.RAG_ANSWER_END));
                streamComplete.complete(null);
            }).onError(throwable -> {
                String errMsg = throwable.getMessage();
                log.error("RAG流式回答生成过程中发生错误", throwable);

                if (errMsg != null && errMsg.contains("AllocationQuota.FreeTierOnly")) {
                    sendSseData(emitter, createErrorResponse("当前模型额度已用完，请联系管理员修复"));
                } else {
                    sendSseData(emitter, createErrorResponse("回答生成失败: " + errMsg));
                }
                streamComplete.completeExceptionally(throwable);
            });

            // 启动流式处理
            tokenStream.start();

            // 阻塞等待流式处理完成，最多等待5分钟
            try {
                streamComplete.get(5, java.util.concurrent.TimeUnit.MINUTES);
            } catch (TimeoutException e) {
                sendSseData(emitter, createErrorResponse("响应超时"));
            } catch (Exception e) {
                log.error("等待LLM流式响应时发生错误", e);
            }

        } catch (Exception e) {
            log.error("RAG流式回答生成发生异常", e);
            sendSseData(emitter, createErrorResponse("回答生成失败: " + e.getMessage()));
        }
    }

    /**
     * 构建流式Agent
     */
    private Agent buildStreamingAgent(StreamingChatModel streamingClient, String sessionId) {
        MessageWindowChatMemory memory = MessageWindowChatMemory.builder()
                .id(sessionId) // 指定 sessionId，会话隔离
                .maxMessages(2) // 保留最近2条（尽量不让上下文污染到模型回答的生成）
                .chatMemoryStore(redisChatMemoryStore) // 存储到 Redis
                .build();

        SystemPrompts sp = promptMapper.selectOne(
                Wrappers.lambdaQuery(SystemPrompts.class)
                        .eq(SystemPrompts::getStatus, Boolean.TRUE)
        );
        if (sp == null) {
            log.error("未找到启用的 SystemPrompt 配置");
            throw new BusinessException("服务暂时不可用，请稍后再试");
        }
        final String systemPrompt = sp.getSystemPrompt();

        memory.add(new SystemMessage(systemPrompt));

        return AiServices.builder(Agent.class)
                .streamingChatModel(streamingClient)
                .chatMemory(memory)
                .build();
    }

}
